#include "slam_controller.h"

#include <visualization_msgs/Marker.h>
#include "SLAM/Location.h"
#include <sensor/IRInfo.h>
#include <fstream>

#define NUM_SENSORS 8

SLAMController::SLAMController() :
	robot_active_(false),
	robot_(0.0f, 0.0f, 0.0f),
	path_()
{
	path_.setOccupancyGrid(mapper_.getOccupancyGrid());
	path_.setTags(&tags_);

	zeroTagHack_ = new TargetTag(0.0f, 0.0f, 0.0f, 0, 0);

//	TargetTag *tag = new TargetTag(2.18f, -0.0f, M_PI-0.1f, 0, 0);
//	tags_.push_back(tag);
//	tag = new TargetTag(2.18f, -0.7f, M_PI-0.1f, 1, 0);
//	tags_.push_back(tag);
//	tag = new TargetTag(1.20f, -0.22f, -M_PI/2.0f-0.1f, 2, 0);
//	tags_.push_back(tag);

	ros::NodeHandle nh("SLAMController");

	last_sensor_readings_ = (SensorRayParameters *) malloc(sizeof(SensorRayParameters) * NUM_SENSORS);
	for (int i = 0; i < NUM_SENSORS; i++) last_sensor_readings_[i].mean = 0.0f;

	// Subscribe to the sensor readings
	sensor_ir_subscriber_ = nh.subscribe("/sensor/ir", 1, &SLAMController::sensor_callback, this);

	sensor_tag_subscriber_ = nh.subscribe("/sensor/tag", 1, &SLAMController::sensor_tag_callback, this);

	// Create publisher to publish visualization calls
	distance_ray_publisher_ = nh.advertise<visualization_msgs::Marker>("/distance_ray_marker", 0);

	// Create the service client for communication with the Tracker
	tracker_client_ = nh.serviceClient<SLAM::Location>("/slam/location");

	// Create the service client for requesting IRInfo from sensor
	sensor_info_client_ = nh.serviceClient<sensor::IRInfo>("/sensor/ir/info");

	// Add the probabilities service
	probabilities_server = nh.advertiseService("/slam/particle_probabilities", &SLAMController::probabilities_request, this);

	path_server = nh.advertiseService("/slam/path", &SLAMController::path_request, this);

	switch_phase_subscriber_ = nh.subscribe("/SLAM/phase", 100, &SLAMController::switch_phase_callback, this);

	recorded_tags_server_ = nh.advertiseService("/SLAM/tags", &SLAMController::recorded_tags_request, this);

	waypoints_server_ = nh.advertiseService("/SLAM/waypoints", &SLAMController::waypoints_request, this);

	getSensorInformation();

	// Start a background spinner that can receive the topic messages
	ros::AsyncSpinner spinner(3);
	spinner.start();

	double lastMapRenderTime = ros::Time::now().toSec();
	// Start a render loop
	ros::Rate loop_rate(10);
	while (ros::ok()) {
		if ((ros::Time::now().toSec() - lastMapRenderTime) > 1.5) {
			// Render the map
			mapper_.render();

//			path_.setTarget(tag);
//			path_.getWaypoints();

			path_.renderInflatedMap();

			path_.renderPath();

			renderTags();

			lastMapRenderTime = ros::Time::now().toSec();
		}

		// Render the robot
		robot_.render();

		loop_rate.sleep();
	}

	ros::waitForShutdown();
}

SLAMController::~SLAMController()
{
	free(last_sensor_readings_);

	deleteAllTags();
}

void SLAMController::deleteAllTags() {
	for (unsigned int i = 0; i < tags_.size(); i++) {
		delete tags_[i];
	}
	tags_.clear();
}

void SLAMController::loadMapDataFromFile() {
	std::ifstream tagsFile("tags.txt");
	std::ifstream mapFile("map.txt");

	int n;
	tagsFile >> n;
	for (int i = 0; i < n; i++) {
		float x, y, angle;
		int tagID, color;
		tagsFile >> x >> y >> angle >> tagID >> color;
		TargetTag *tag = new TargetTag(x, y, angle, tagID, color);
		tags_.push_back(tag);
	}

	mapper_.getOccupancyGrid()->readFromFile(mapFile);
}

void SLAMController::saveMapDataToFile() {
	std::ofstream tagsFile("tags.txt");
	std::ofstream mapFile("map.txt");

	if (!tagsFile.is_open() || !mapFile.is_open()) {
		ROS_ERROR("Could not open tag or map file.");
	}

	tagsFile << tags_.size() << std::endl;
	for (unsigned int i = 0; i < tags_.size(); i++) {
		TargetTag *tag = tags_[i];
		tagsFile << tag->getX() << std::endl;
		tagsFile << tag->getY() << std::endl;
		tagsFile << tag->getAngle() << std::endl;
		tagsFile << tag->getTagID() << std::endl;
		tagsFile << tag->getColor() << std::endl;
	}

	mapper_.getOccupancyGrid()->writeToFile(mapFile);

	tagsFile.close();
	mapFile.close();
}

void SLAMController::sensor_callback(const sensor::IR::ConstPtr& msg)
{
	if (!robot_active_) return;

	// Get the location from the Tracker
	SLAM::Location srv;
	srv.request.resample = true;
	if (tracker_client_.call(srv)) {
		path_.setCurrentPosition(srv.response.x, srv.response.y, srv.response.angle);

		// Update robot position
		robot_.setX(srv.response.x);
		robot_.setY(srv.response.y);
		robot_.setAngle(srv.response.angle);

		vec2 robotLocation;
		robotLocation.x = robot_.getX();
		robotLocation.y = robot_.getY();
		mapper_.addHitWithDistribution(robotLocation, 0.10f, -0.5f);

		// Ignore sensor readings when robot not aligned
		// to the wall.
		float leftFront	= msg -> dists[0];
		float leftBack 	= msg -> dists[1];
		float rightFront  = msg -> dists[5];
		float rightBack   = msg -> dists[4];

		// Normalize values.
		if (leftFront < 0.01f || leftFront > 0.30f)
			leftFront = 0.30f;
		if (leftBack < 0.01f || leftBack > 0.30f)
			leftBack = leftFront;
		if (rightFront < 0.01f || rightFront > 0.30f)
			rightFront = 0.30f;
		if (rightBack < 0.01f || rightBack > 0.30f)
			rightBack = rightFront;

		bool updateMap = true;

		float sideSensorDist = 0.075f;
		float leftAngle = tanf(fabs(leftFront - leftBack) / sideSensorDist);
		float rightAngle = tanf(fabs(rightFront - rightBack) / sideSensorDist);

		updateMap = leftAngle < 0.17f || rightAngle < 0.17f;

		float modAngle = fmodf(robot_.getAngle(), M_PI / 2.0f);
		if (updateMap && fabsf(modAngle) < 0.17f) {
			robot_.setAngle(robot_.getAngle() - modAngle * 0.2f);
		}

		vec2 from;
		float cosAngle = cosf(robot_.getAngle());
		float sinAngle = sinf(robot_.getAngle());

		// For each sensor, translate the sensor ray relative to the robot's
		// position and enter them into the mapper.
		for (int i = 0; i < NUM_SENSORS; i++) {
			float distance = msg->dists[i];
			last_sensor_readings_[i].mean = distance;

			if (distance < 0.01f) {
				last_sensor_readings_[i].mean = distance;
				continue;
			}
			if (i == 2 || i == 3) continue; // don't use rear long range sensors

			// Calculate the origin of the sensor ray
			from.x = robot_.getX() + sensor_infos_[i].x * cosAngle - sensor_infos_[i].y * sinAngle;
			from.y = robot_.getY() + sensor_infos_[i].x * sinAngle + sensor_infos_[i].y * cosAngle;

			// Direction the ray travels in absolute angle
			float rayAngle = robot_.getAngle() + sensor_infos_[i].angle;

			float std_dev = 0.01f;
			float height = 1.0f;
			if (updateMap) {
				mapper_.addSensorRay(from, rayAngle, distance, std_dev, height);
			}
			// Save parameters
			last_sensor_readings_[i].std_dev = std_dev;
			last_sensor_readings_[i].height = height;
			//mapper_.addHitWithDistribution(to, 0.02f, 0.3f);
			//renderRay(from.x, from.y, to.x, to.y);
		}

		path_.update();
	} else {
		ROS_ERROR("SLAMController: Failed to get Location from Tracker.");
	}
}

void SLAMController::sensor_tag_callback(const sensor::tag::ConstPtr& msg)
{
	if (!robot_active_) return;

	SLAM::Location loc;
	loc.request.resample = false;
	if (tracker_client_.call(loc)) {
		// Calculate tag location
		float tag_x = loc.response.x - 0.10f * cosf(loc.response.angle);
		float tag_y = loc.response.y - 0.10f * sinf(loc.response.angle);

		bool foundMatchingTag = false;

		// Inspect tags to see if the received tag is closer than 30cm to any other tag
		for (unsigned int i = 0; i < tags_.size(); i++) {
			// Calculate distance between the received tag and the current tag
			float dx = tag_x - tags_[i]->getX();
			float dy = tag_y - tags_[i]->getY();
			float distance = sqrtf(dx * dx + dy * dy);

			if (distance < 0.30f) {
				foundMatchingTag = true;
				break;
			}
		}

		if (!foundMatchingTag) {
			TargetTag *tag = new TargetTag(tag_x, tag_y, msg->angle - loc.response.angle, 0, 0);
			tags_.push_back(tag);
		}
	} else {
		ROS_ERROR("SLAMController: Failed to get Location from Tracker.");
	}
}

bool SLAMController::probabilities_request(SLAM::Probabilities::Request& req, SLAM::Probabilities::Response& res)
{
	int num_received = req.particles.size() / 3;
	std::vector<uint8_t> probabilities(num_received);
	for (int i = 0; i < num_received; i++) {
		float x = req.particles[i * 3];
		float y = req.particles[i * 3 + 1];
		float theta = req.particles[i * 3 + 2];

		vec2 from;
		float cosAngle = cosf(theta);
		float sinAngle = sinf(theta);

		// For each sensor, translate the sensor ray relative to the robot's
		// position and enter them into the mapper.
		float total_probability = 0;
		int sensor_count = 0;
		for (int i = 0; i < NUM_SENSORS; i++) {
			SensorRayParameters params = last_sensor_readings_[i];
			if (params.mean < 0.01f) continue;
			if (i == 2 || i == 3) continue; // don't use rear long range sensors

			// Calculate the origin of the sensor ray
			from.x = x + sensor_infos_[i].x * cosAngle - sensor_infos_[i].y * sinAngle;
			from.y = y + sensor_infos_[i].x * sinAngle + sensor_infos_[i].y * cosAngle;

			// Direction the ray travels in absolute angle
			float rayAngle = theta + sensor_infos_[i].angle;

			float prob = mapper_.probabilityOfSensorRay(from, rayAngle, params.mean, params.std_dev, params.height);
			total_probability += prob;
			sensor_count++;
		}
		// Push the average of the probabilities of the sensor readings
		res.probabilities.push_back((uint8_t) ((total_probability / sensor_count) * 255.0f));
	}

	return num_received > 0;
}

bool SLAMController::path_request(SLAM::Directions::Request& req, SLAM::Directions::Response& res)
{
	SLAM::Location loc;
	loc.request.resample = false;
	if (tracker_client_.call(loc)) {
		path_.setCurrentPosition(loc.response.x, loc.response.y, loc.response.angle);
		path_.setTarget(NULL);
		for (unsigned int i = 0; i < tags_.size(); i++) {
			if (tags_[i]->getTagID() == req.tagID) {
				path_.setTarget(tags_[i]);
				break;
			}
		}

		if (path_.getNumWaypoints() == 0) {
			res.distance = 0;
			res.direction = 0;
		} else {
			Waypoint p = path_.getFirstWaypoint();
			float dx = loc.response.x - p.x;
			float dy = loc.response.y - p.y;
			float dist = sqrtf(dx * dx + dy * dy);

			res.distance = dist;
			res.direction = atan2f(dx, dy) - loc.response.angle;
		}
	} else {
		ROS_ERROR("Mapper: Failed to get Location from Tracker.");
	}

	return true;
}

void SLAMController::switch_phase_callback(const SLAM::SwitchPhase::ConstPtr& msg) {
	/*
		0 = Do nothing
		1 = Begin exploration
		2 = Exploration finished
		3 = Start fetching
		4 = Fetching finished
	*/

	switch (msg->phase) {
		case 0:
			robot_active_ = false;
			break;
		case 1:
			robot_.setX(0.0f);
			robot_.setY(0.0f);
			robot_.setAngle(0.0f);
			path_.setCurrentPosition(0.0f, 0.0f, 0.0f);

			robot_active_ = true;
			break;
		case 2:
			robot_active_ = false;
			saveMapDataToFile();
			break;
		case 3:
			deleteAllTags();
			loadMapDataFromFile();

			robot_.setX(0.0f);
			robot_.setY(0.0f);
			robot_.setAngle(0.0f);
			path_.setCurrentPosition(0.0f, 0.0f, 0.0f);

			robot_active_ = true;
			break;
		case 4:
			robot_active_ = false;
			break;
	}
}

bool SLAMController::recorded_tags_request(SLAM::RecordedTags::Request& req, SLAM::RecordedTags::Response& res) {
	for (unsigned int i = 0; i < tags_.size(); i++) {
		TargetTag *tag = tags_[i];

		SLAM::Tag tagMsg;
		tagMsg.index = i;
		tagMsg.x = tag->getX();
		tagMsg.y = tag->getY();

		res.tags.push_back(tagMsg);
	}

	return true;
}

bool SLAMController::waypoints_request(SLAM::Waypoints::Request& req, SLAM::Waypoints::Response& res) {
	SLAM::Location loc;
	loc.request.resample = false;
	if (tracker_client_.call(loc)) {
		path_.setCurrentPosition(loc.response.x, loc.response.y, loc.response.angle);
		if (req.index == -1) {
			path_.setTarget(zeroTagHack_);
		} else {
			if (req.index < 0 || req.index >= (int) tags_.size()) {
				ROS_ERROR("Waypoints request: Requested tag index out of bounds %d", req.index);
				return false;
			}

			path_.setTarget(tags_[req.index]);
		}

		std::vector<Waypoint> waypoints = path_.getWaypoints();

		for (unsigned int i = 0; i < waypoints.size(); i++) {
			SLAM::Waypoint wp;
			wp.x = waypoints[i].x;
			wp.y = waypoints[i].y;
			res.waypoints.push_back(wp);
		}
	} else {
		ROS_ERROR("Mapper: Failed to get Location from Tracker.");
	}

	return true;
}

void SLAMController::renderRay(float from_x, float from_y, float to_x, float to_y)
{
	visualization_msgs::Marker marker;

	marker.header.frame_id = "/floor";
	marker.header.stamp = ros::Time();

	marker.ns = "distance_ray_marker";
	marker.id = random();
	marker.lifetime = ros::Duration(0.1f);

	marker.type = visualization_msgs::Marker::LINE_LIST;
	marker.action = visualization_msgs::Marker::ADD;

	marker.pose.position.x = 0.0f;
	marker.pose.position.y = 0.0f;
	marker.pose.position.z = 0.0f;

	marker.pose.orientation.x = 0.0f;
	marker.pose.orientation.y = 0.0f;
	marker.pose.orientation.z = 0.0f;
	marker.pose.orientation.w = 1.0f;

	marker.scale.x = 0.01f;
	marker.scale.y = 1.0f;
	marker.scale.z = 1.0f;

	marker.color.r = 1.0f;
	marker.color.g = 0.5f;
	marker.color.b = 0.0f;
	marker.color.a = 1.0f;

	std::vector<geometry_msgs::Point> points(2);
	points[0].x = from_x;
	points[0].y = from_y;
	points[0].z = 0.0f;
	points[1].x = to_x;
	points[1].y = to_y;
	points[1].z = 0.0f;
	marker.points = points;

	// Publish the ray
	distance_ray_publisher_.publish(marker);
}

void SLAMController::renderTags()
{
	visualization_msgs::Marker marker;

	marker.header.frame_id = "/floor";
	marker.header.stamp = ros::Time();

	marker.ns = "distance_ray_marker";
	marker.id = random();
	marker.lifetime = ros::Duration(1.5f);

	marker.type = visualization_msgs::Marker::LINE_LIST;
	marker.action = visualization_msgs::Marker::ADD;

	marker.pose.position.x = 0.0f;
	marker.pose.position.y = 0.0f;
	marker.pose.position.z = 0.01f;

	marker.pose.orientation.x = 0.0f;
	marker.pose.orientation.y = 0.0f;
	marker.pose.orientation.z = 0.0f;
	marker.pose.orientation.w = 1.0f;

	std::vector<geometry_msgs::Point> points(tags_.size() * 2);

	for (unsigned int i = 0; i < tags_.size(); i++) {
		TargetTag *tag = tags_[i];
		points[i * 2].x = tag->getX();
		points[i * 2].y = tag->getY();
		points[i * 2].z = 0.0f;
		points[i * 2 + 1].x = tag->getX() + 0.2f * cosf(tag->getAngle());
		points[i * 2 + 1].y = tag->getY() + 0.2f * sinf(tag->getAngle());
		points[i * 2 + 1].z = 0.0f;
	}

	marker.points = points;

	marker.scale.x = 0.02f;
	marker.scale.y = 0.1f;
	marker.scale.z = 0.1f;

	marker.color.r = 0.0f;
	marker.color.g = 1.0f;
	marker.color.b = 0.0f;
	marker.color.a = 1.0f;

	// Publish the ray
	distance_ray_publisher_.publish(marker);
}

void SLAMController::getSensorInformation()
{
	// Wait for the sensor info service to become available with
	// a timeout of 10 seconds.
	sensor_info_client_.waitForExistence(ros::Duration(10.0f));

	for (int i = 0; i < NUM_SENSORS; i++) {
		sensor::IRInfo infoSrv;
		infoSrv.request.sensor_id = i;

		if (sensor_info_client_.call(infoSrv)) {
			IRSensorInfo sensorInfo;

			sensorInfo.x = infoSrv.response.x;
			sensorInfo.y = infoSrv.response.y;
			sensorInfo.z = infoSrv.response.z;
			sensorInfo.angle = infoSrv.response.angle;

			sensor_infos_.push_back(sensorInfo);
		} else {
			ROS_ERROR("SLAMController: Failed call to sensor info client for sensor %d", i);
		}
	}
}

int main(int argc, char **argv)
{
	ros::init(argc, argv, "mapper");

	new SLAMController();

	return 0;
}
